import argparse
import os
from predictor import ImagePredictor
from config import Config

def main():
    parser = argparse.ArgumentParser(description='图像识别系统')
    subparsers = parser.add_subparsers(dest='command', help='可用命令')
    
    # 训练命令
    train_parser = subparsers.add_parser('train', help='训练模型')
    train_parser.add_argument('--data-dir', type=str, help='训练数据目录')
    train_parser.add_argument('--config', type=str, default='config.yml', help='配置文件路径')
    
    # 预测命令
    predict_parser = subparsers.add_parser('predict', help='预测图片')
    predict_parser.add_argument('image_path', type=str, help='要预测的图片路径')
    predict_parser.add_argument('--model', type=str, required=True, help='模型路径')
    predict_parser.add_argument('--classes', type=str, required=True, help='类别名称文件路径')
    predict_parser.add_argument('--top-k', type=int, default=3, help='返回前k个预测结果')
    
    args = parser.parse_args()
    
    if args.command == 'train':
        # 加载配置
        config = Config(args.config)
        # 这里调用训练流程
        from main import main as train_main
        train_main(config)
        
    elif args.command == 'predict':
        if not os.path.exists(args.image_path):
            print(f"错误：图片不存在 - {args.image_path}")
            return
            
        predictor = ImagePredictor(args.model, args.classes)
        results = predictor.predict(args.image_path, args.top_k)
        
        print("\n预测结果：")
        for i, result in enumerate(results, 1):
            print(f"{i}. {result['class']}: {result['probability']:.2%}")

if __name__ == '__main__':
    main() 